{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|default_exp models.RNN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RNNs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ">These are RNN, LSTM and GRU PyTorch implementations created by Ignacio Oguiza - oguiza@timeseriesAI.co"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "from tsai.imports import *\n",
    "from tsai.models.layers import *\n",
    "from tsai.models.utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class _RNN_Base(Module):\n",
    "    def __init__(self, c_in, c_out, hidden_size=100, n_layers=1, bias=True, rnn_dropout=0, bidirectional=False, fc_dropout=0., init_weights=True):\n",
    "        self.rnn = self._cell(c_in, hidden_size, num_layers=n_layers, bias=bias, batch_first=True, dropout=rnn_dropout, \n",
    "                              bidirectional=bidirectional)\n",
    "        self.dropout = nn.Dropout(fc_dropout) if fc_dropout else nn.Identity()\n",
    "        self.fc = nn.Linear(hidden_size * (1 + bidirectional), c_out)\n",
    "        if init_weights: self.apply(self._weights_init)\n",
    "\n",
    "    def forward(self, x): \n",
    "        x = x.transpose(2,1)    # [batch_size x n_vars x seq_len] --> [batch_size x seq_len x n_vars]\n",
    "        output, _ = self.rnn(x) # output from all sequence steps: [batch_size x seq_len x hidden_size * (1 + bidirectional)]\n",
    "        output = output[:, -1]  # output from last sequence step : [batch_size x hidden_size * (1 + bidirectional)]\n",
    "        output = self.fc(self.dropout(output))\n",
    "        return output\n",
    "    \n",
    "    def _weights_init(self, m): \n",
    "        # same initialization as keras. Adapted from the initialization developed \n",
    "        # by JUN KODA (https://www.kaggle.com/junkoda) in this notebook\n",
    "        # https://www.kaggle.com/junkoda/pytorch-lstm-with-tensorflow-like-initialization\n",
    "        for name, params in m.named_parameters():\n",
    "            if \"weight_ih\" in name: \n",
    "                nn.init.xavier_normal_(params)\n",
    "            elif 'weight_hh' in name: \n",
    "                nn.init.orthogonal_(params)\n",
    "            elif 'bias_ih' in name:\n",
    "                params.data.fill_(0)\n",
    "                # Set forget-gate bias to 1\n",
    "                n = params.size(0)\n",
    "                params.data[(n // 4):(n // 2)].fill_(1)\n",
    "            elif 'bias_hh' in name:\n",
    "                params.data.fill_(0)\n",
    "        \n",
    "class RNN(_RNN_Base):\n",
    "    _cell = nn.RNN\n",
    "    \n",
    "class LSTM(_RNN_Base):\n",
    "    _cell = nn.LSTM\n",
    "    \n",
    "class GRU(_RNN_Base):\n",
    "    _cell = nn.GRU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 16\n",
    "c_in = 3\n",
    "seq_len = 12\n",
    "c_out = 2\n",
    "xb = torch.rand(bs, c_in, seq_len)\n",
    "test_eq(RNN(c_in, c_out, hidden_size=100, n_layers=2, bias=True, rnn_dropout=0.2, bidirectional=True, fc_dropout=0.5)(xb).shape, [bs, c_out])\n",
    "test_eq(RNN(c_in, c_out)(xb).shape, [bs, c_out])\n",
    "test_eq(RNN(c_in, c_out, hidden_size=100, n_layers=2, bias=True, rnn_dropout=0.2, bidirectional=True, fc_dropout=0.5)(xb).shape, [bs, c_out])\n",
    "test_eq(LSTM(c_in, c_out)(xb).shape, [bs, c_out])\n",
    "test_eq(GRU(c_in, c_out)(xb).shape, [bs, c_out])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tsai.basics import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<style>\n",
       "    /* Turns off some styling */\n",
       "    progress {\n",
       "        /* gets rid of default border in Firefox and Opera. */\n",
       "        border: none;\n",
       "        /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "        background-size: auto;\n",
       "    }\n",
       "    progress:not([value]), progress:not([value])::-webkit-progress-bar {\n",
       "        background: repeating-linear-gradient(45deg, #7e7e7e, #7e7e7e 10px, #5c5c5c 10px, #5c5c5c 20px);\n",
       "    }\n",
       "    .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "        background: #F44336;\n",
       "    }\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>epoch</th>\n",
       "      <th>train_loss</th>\n",
       "      <th>valid_loss</th>\n",
       "      <th>accuracy</th>\n",
       "      <th>time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>1.743440</td>\n",
       "      <td>1.633068</td>\n",
       "      <td>0.361111</td>\n",
       "      <td>00:01</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dsid = 'NATOPS' \n",
    "bs = 16\n",
    "X, y, splits = get_UCR_data(dsid, return_split=False)\n",
    "tfms  = [None, [TSCategorize()]]\n",
    "dsets = TSDatasets(X, y, tfms=tfms, splits=splits)\n",
    "dls   = TSDataLoaders.from_dsets(dsets.train, dsets.valid, bs=bs, num_workers=0, shuffle=False)\n",
    "model = LSTM(dls.vars, dls.c)\n",
    "learn = Learner(dls, model,  metrics=accuracy)\n",
    "learn.fit_one_cycle(1, 3e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RNN(\n",
      "  (rnn): RNN(3, 100, num_layers=2, batch_first=True, dropout=0.5, bidirectional=True)\n",
      "  (dropout): Dropout(p=0.5, inplace=False)\n",
      "  (fc): Linear(in_features=200, out_features=2, bias=True)\n",
      ")\n",
      "81802\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 2])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m = RNN(c_in, c_out, hidden_size=100,n_layers=2,bidirectional=True,rnn_dropout=.5,fc_dropout=.5)\n",
    "print(m)\n",
    "print(count_parameters(m))\n",
    "m(xb).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LSTM(\n",
      "  (rnn): LSTM(3, 100, num_layers=2, batch_first=True, dropout=0.5, bidirectional=True)\n",
      "  (dropout): Dropout(p=0.5, inplace=False)\n",
      "  (fc): Linear(in_features=200, out_features=2, bias=True)\n",
      ")\n",
      "326002\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 2])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m = LSTM(c_in, c_out, hidden_size=100,n_layers=2,bidirectional=True,rnn_dropout=.5,fc_dropout=.5)\n",
    "print(m)\n",
    "print(count_parameters(m))\n",
    "m(xb).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GRU(\n",
      "  (rnn): GRU(3, 100, num_layers=2, batch_first=True, dropout=0.5, bidirectional=True)\n",
      "  (dropout): Dropout(p=0.5, inplace=False)\n",
      "  (fc): Linear(in_features=200, out_features=2, bias=True)\n",
      ")\n",
      "244602\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 2])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m = GRU(c_in, c_out, hidden_size=100,n_layers=2,bidirectional=True,rnn_dropout=.5,fc_dropout=.5)\n",
    "print(m)\n",
    "print(count_parameters(m))\n",
    "m(xb).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Converting a model to TorchScript"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.0287, -0.0105]], grad_fn=<AddmmBackward0>)\n"
     ]
    }
   ],
   "source": [
    "model = LSTM(c_in, c_out, hidden_size=100, n_layers=2, bidirectional=True, rnn_dropout=.5, fc_dropout=.5)\n",
    "model.eval()\n",
    "inp = torch.rand(1, c_in, 50)\n",
    "output = model(inp)\n",
    "print(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tracing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LSTM(\n",
      "  original_name=LSTM\n",
      "  (rnn): LSTM(original_name=LSTM)\n",
      "  (dropout): Dropout(original_name=Dropout)\n",
      "  (fc): Linear(original_name=Linear)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# save to gpu, cpu or both\n",
    "traced_cpu = torch.jit.trace(model.cpu(), inp)\n",
    "print(traced_cpu)\n",
    "torch.jit.save(traced_cpu, \"cpu.pt\")\n",
    "\n",
    "# load cpu or gpu model\n",
    "traced_cpu = torch.jit.load(\"cpu.pt\")\n",
    "test_eq(traced_cpu(inp), output)\n",
    "\n",
    "!rm \"cpu.pt\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scripting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "RecursiveScriptModule(\n",
      "  original_name=LSTM\n",
      "  (rnn): RecursiveScriptModule(original_name=LSTM)\n",
      "  (dropout): RecursiveScriptModule(original_name=Dropout)\n",
      "  (fc): RecursiveScriptModule(original_name=Linear)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# save to gpu, cpu or both\n",
    "scripted_cpu = torch.jit.script(model.cpu())\n",
    "print(scripted_cpu)\n",
    "torch.jit.save(scripted_cpu, \"cpu.pt\")\n",
    "\n",
    "# load cpu or gpu model\n",
    "scripted_cpu = torch.jit.load(\"cpu.pt\")\n",
    "test_eq(scripted_cpu(inp), output)\n",
    "\n",
    "!rm \"cpu.pt\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Converting a model to ONNX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#| onnx\n",
    "```python\n",
    "import onnx\n",
    "\n",
    "# Export the model\n",
    "torch.onnx.export(model.cpu(),               # model being run\n",
    "                  inp,                       # model input (or a tuple for multiple inputs)\n",
    "                  \"cpu.onnx\",                # where to save the model (can be a file or file-like object)\n",
    "                  export_params=True,        # store the trained parameter weights inside the model file\n",
    "                  verbose=False,\n",
    "                  opset_version=13,          # the ONNX version to export the model to\n",
    "                  do_constant_folding=True,  # whether to execute constant folding for optimization\n",
    "                  input_names = ['input'],   # the model's input names\n",
    "                  output_names = ['output'], # the model's output names\n",
    "                  dynamic_axes={\n",
    "                      'input'  : {0 : 'batch_size'}, \n",
    "                      'output' : {0 : 'batch_size'}} # variable length axes\n",
    "                 )\n",
    "\n",
    "# Load the model and check it's ok\n",
    "onnx_model = onnx.load(\"cpu.onnx\")\n",
    "onnx.checker.check_model(onnx_model)\n",
    "\n",
    "# You can ignore the WARNINGS below\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#| onnx\n",
    "```python\n",
    "import onnxruntime as ort\n",
    "\n",
    "ort_sess = ort.InferenceSession('cpu.onnx')\n",
    "out = ort_sess.run(None, {'input': inp.numpy()})\n",
    "\n",
    "# input & output names\n",
    "input_name = ort_sess.get_inputs()[0].name\n",
    "output_name = ort_sess.get_outputs()[0].name\n",
    "\n",
    "# input dimensions\n",
    "input_dims = ort_sess.get_inputs()[0].shape\n",
    "print(input_name, output_name, input_dims)\n",
    "\n",
    "test_close(out, output.detach().numpy())\n",
    "!rm \"cpu.onnx\"\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export -"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": "IPython.notebook.save_checkpoint();",
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/nacho/notebooks/tsai/nbs/042_models.RNN.ipynb saved at 2022-12-10 11:04:07\n",
      "Correct notebook to script conversion! 😃\n",
      "Saturday 10/12/22 11:04:10 CET\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" autoplay=\"autoplay\">\n",
       "                    <source src=\"data:audio/wav;base64,UklGRvQHAABXQVZFZm10IBAAAAABAAEAECcAACBOAAACABAAZGF0YdAHAAAAAPF/iPh/gOoOon6w6ayCoR2ZeyfbjobxK+F2Hs0XjKc5i3DGvzaTlEaraE+zz5uLUl9f46fHpWJdxVSrnfmw8mYEScqUP70cb0Q8X41uysJ1si6Eh1jYzXp9IE2DzOYsftYRyoCY9dJ/8QICgIcEun8D9PmAaBPlfT7lq4MFIlh61tYPiCswIHX+yBaOqT1QbuW7qpVQSv9lu6+xnvRVSlyopAypbGBTUdSalrSTaUBFYpInwUpxOzhti5TOdndyKhCGrdwAfBUcXIJB69p+Vw1egB76+n9q/h6ADglbf4LvnIHfF/981ODThF4m8HiS0riJVjQ6c+/EOZCYQfJrGrhBmPVNMmNArLKhQlkXWYqhbaxXY8ZNHphLuBJsZUEckCTFVHMgNKGJytIDeSUmw4QN4Qx9pReTgb3vYX/TCBuApf75f+P5Y4CRDdN+B+tngk8c8nt03CKGqipgd13OhotwOC5x9MCAknFFcmlmtPmagFFFYOCo0qRzXMhVi57pryNmIEqJlRi8bm52PfuNM8k4dfQv+4cO12l6zCGdg3jl730uE/KAPvS+f0wEAoAsA89/XfXQgBESIn6S5luDtiC8eh/YmIfpLqt1OMp5jXg8/24MveqUNUnPZsqw0Z3yVDldnaUOqIZfXlKrm36zzWhjRhaT+r+ncHI5/otUzfd2uSt7hl/bqXtoHaCC6+mqfrAOeoDD+PJ/xf8RgLMHfH/b8GeBihZIfSXidoQSJWB52NM1iRkzz3MkxpKPbUCrbDu5d5fgTAxkSK3JoEhYD1p2omere2LZTuqYLbdWa49Cx5Dww7tyXDUnioXRkHhwJyKFvd/AfPoYy4Fl7j1/LQorgEr9/X89+0qAOAwAf13sJoL8Gkd8wt25hWIp3Heez/eKODfPcSPCzpFNRDVqf7UlmnNQKGHgqd+jgVvJVm2f265QZTpLS5byur1tpT6ajvrHq3Q2MXWIxtUCehoj8YMk5LB9hRQegeTypn+nBQWA0QHgf7f2q4C5EFt+5ucOg2YfHXtq2SSHpS0ydnTL4IxFO6pvNb4ulBdInWfcsfSc7VMmXpSmE6eeXmZThJxpsgRohEfOk86+AHCoOpOMFsx1dv8s6oYT2k17uR7ngpXod34IEJqAaPfnfyABCIBZBpl/NPI2gTQVjX134x2ExSPMeR7VtYjZMWJ0W8ftjkA/YW1durCWykvjZFKu4p9LVwVbZKNkqpxh6U+6mRC2mGq2Q3SRvsIgcpc2sIpD0Bp4uiiFhW3ecXxOGgaCDe0Vf4cLPoDv+/5/mfw1gN4KKX+17emBqBmYfBHfVYUZKFR44NBtiv41bHJUwx+RJkP1apu2VJlkTwli4qrwoo1ax1dToNCtemRSTBGXz7kJbdM/PY/Dxht0dTLziH7Ul3loJEiE0uJsfdsVTYGL8Yt/AgcMgHYA7X8S+IqAYA+QfjzpxIIVHnp7tdqzhmAstXaxzEqMETpScGC/dJP3Rmdo8LIZnOVSEF+Opxumsl1sVF+dVrE5Z6NIiZSkvVdv2zsqjdnK8HVDLlyHyNjuegogM4NA5z9+YRG9gA722H97AgOA/gSyf43zCIHdE899yuTIg3ciNXpm1jmImTDwdJPITI4RPhRugbvslbFKt2Vfr/6eTFb4W1WkY6m6YPdQjJr2tNZp3EQlko7BgXHRNz2LAc+gdwMq7IUf3R58ohtFgrbr6n7hDFWAlPr8f/T9I4CECU9/De+vgVQY5nxh4POEzybJeCTS5YnCNAZzhsRzkP1Bsmu4t4aYU07nYuerA6KWWcJYO6HHrKJjaE3Zl624UWz/QOOPjcWHc7QzdIk40yl5tCWjhIDhJX0xF4CBMvBsf10IF4Ac//Z/bPlsgAcOwn6S6n6CwxzUewLcRoYaKzV38M23i9o493CNwL6S1UUuaQe0QpvbUfdfiqglpcRccFU+nkWwambASUiVfLyqbg49xY2eyWh1hy/Sh37XjHpaIYKD7OUEfrgS5IC09MV/1gMBgKMDyH/n9N6AhhINfh7mdoMoIZt6r9fAh1cvfHXNya6N4DzDbqi8K5WWSYlmbbAdnkpV6FxJpWSo1V8DUmGb3rMRaQBG2JJgwN9wCDnNi8HNI3dKK1aG0dvHe/UciIJf6rt+Og5wgDn59X9P/xWAKQhxf2XweYH+FjB9suGVhIMlOnlo02GJhTOdc7vFyo/TQGxs2Li7lz9NwmPurBihnVi7WSWiwKvGYntOpJiOt5drKUKMkFnE8HLxNPmJ9NG4eP8mAYUv4Np8hhi3gdruSX+3CSWAwP38f8f6UoCuDPF+6Os8gnAbKnxQ3d2F0imydzDPKIuiN5lxu8EKkrFE82kftW2az1DbYImpMqTUW3FWIJ83r5hl2koJlla7+m0+PmSOZcjcdMgwS4g11iZ6qCLUg5jkxn0QFA6BWvOvfzEFBIBHAtp/Qfa3gC4RSH5y5yeD2B/8evnYS4cULgR2CMsUja47cG/QvW6UeEhXZ3+xP51GVNVdP6Zpp+1eDFM5nMeySWghR4+TNL85cD46YIyCzKJ2kCzEhoTabXtGHs+CCemJfpMPjoDe9+t/qQALgM8Gj3++8UaBqRV2fQTjO4Q3JKd5r9TgiEYyMHTxxiWPpz8jbfq585YpTJpk960xoKFXsVoTo7yq6GGMTw==\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#|eval: false\n",
    "#|hide\n",
    "from tsai.export import get_nb_name; nb_name = get_nb_name(locals())\n",
    "from tsai.imports import create_scripts; create_scripts(nb_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
